Skip to content

ENH: fix is_jax_array under torch.compile - #464

Closed
lucascolley wants to merge 1 commit into
data-apis:mainfrom
lucascolley:is-jax-torch-compile
Closed

ENH: fix is_jax_array under torch.compile#464
lucascolley wants to merge 1 commit into
data-apis:mainfrom
lucascolley:is-jax-torch-compile

Conversation

@lucascolley

@lucascolley lucascolley commented Aug 16, 2026

Copy link
Copy Markdown
Member

This fixes errors like the following:

E               torch._dynamo.exc.UserError:
GetAttrVariable(TupleVariable(length=0), dtype) has no type
E               For more information about this error, see:
https://pytorch.org/docs/main/generated/exportdb/index.html#unknown-python-type
E
E               from user code:
E                  File
"/Users/lucascolley/ghq/github.com/data-apis/array-api-extra/src/array_api_extra/_lib/_helpers.py",
line 612, in inner
E                   res = func(*args, **kwargs)  # pyright:
ignore[reportCallIssue]
E                 File
"/Users/lucascolley/ghq/github.com/data-apis/array-api-extra/tests/main/test_at.py",
line 47, in at_op
E                   return meth(y, copy=copy, xp=xp)
E                 File
"/Users/lucascolley/ghq/github.com/data-apis/array-api-extra/src/array_api_extra/_at.py",
line 374, in add
E                   return self._op(_AtOp.ADD, operator.iadd,
operator.add, y, copy=copy, xp=xp)  # pyright:
ignore[reportUnknownArgumentType]
E                 File
"/Users/lucascolley/ghq/github.com/data-apis/array-api-extra/src/array_api_extra/_at.py",
line 289, in _op
E                   (_compat.is_dask_array(idx) or
_compat.is_jax_array(idx))
E                 File
"/Users/lucascolley/ghq/github.com/data-apis/array-api-extra/.pixi/envs/tests-backends/lib/python3.14/site-packages/array_api_compat/common/_helpers.py",
line 253, in is_jax_array
E                   or _is_jax_zero_gradient_array(x)
E                 File
"/Users/lucascolley/ghq/github.com/data-apis/array-api-extra/.pixi/envs/tests-backends/lib/python3.14/site-packages/array_api_compat/common/_helpers.py",
line 84, in _is_jax_zero_gradient_array
E                   cls = cast(Hashable, type(dtype))

cc @crusaderky for vis

@lucascolley
BUG: fix is_jax_array under torch.compile
1f66a47 
This fixes errors like the following:
```
E               torch._dynamo.exc.UserError:
GetAttrVariable(TupleVariable(length=0), dtype) has no type
E               For more information about this error, see:
https://pytorch.org/docs/main/generated/exportdb/index.html#unknown-python-type
E
E               from user code:
E                  File
"/Users/lucascolley/ghq/github.com/data-apis/array-api-extra/src/array_api_extra/_lib/_helpers.py",
line 612, in inner
E                   res = func(*args, **kwargs)  # pyright:
ignore[reportCallIssue]
E                 File
"/Users/lucascolley/ghq/github.com/data-apis/array-api-extra/tests/main/test_at.py",
line 47, in at_op
E                   return meth(y, copy=copy, xp=xp)
E                 File
"/Users/lucascolley/ghq/github.com/data-apis/array-api-extra/src/array_api_extra/_at.py",
line 374, in add
E                   return self._op(_AtOp.ADD, operator.iadd,
operator.add, y, copy=copy, xp=xp)  # pyright:
ignore[reportUnknownArgumentType]
E                 File
"/Users/lucascolley/ghq/github.com/data-apis/array-api-extra/src/array_api_extra/_at.py",
line 289, in _op
E                   (_compat.is_dask_array(idx) or
_compat.is_jax_array(idx))
E                 File
"/Users/lucascolley/ghq/github.com/data-apis/array-api-extra/.pixi/envs/tests-backends/lib/python3.14/site-packages/array_api_compat/common/_helpers.py",
line 253, in is_jax_array
E                   or _is_jax_zero_gradient_array(x)
E                 File
"/Users/lucascolley/ghq/github.com/data-apis/array-api-extra/.pixi/envs/tests-backends/lib/python3.14/site-packages/array_api_compat/common/_helpers.py",
line 84, in _is_jax_zero_gradient_array
E                   cls = cast(Hashable, type(dtype))
```
@lucascolley
lucascolley requested a review from ev-br August 16, 2026 14:03
@ev-br

ev-br commented Aug 16, 2026

Copy link
Copy Markdown
Member

Is there an MWE, ideally as a test or as an example for the added comment?

attrs on empty tuples is a tad cryptic, and I'm not sure how to trigger it---or check if the whole _is_zero_grad_array can be removed, as the comment above indicates will be possible "at some point".

@lucascolley

lucascolley commented Aug 16, 2026

Copy link
Copy Markdown
Member Author

calling is_jax_array on a torch array while running under torch.compile(fullgraph=True) should reproduce.

@ev-br

ev-br commented Aug 16, 2026

Copy link
Copy Markdown
Member

Great. Add an otherwise failing code snippet to the comment?

@lucascolley lucascolley mentioned this pull request Aug 17, 2026
Draft
@ev-br

ev-br commented Aug 17, 2026

Copy link
Copy Markdown
Member

CI failures in test_trace are unrelated, data-apis/array-api-tests#455

calling is_jax_array on a torch array while running under torch.compile(fullgraph=True) should reproduce.

Does not repro locally:

In [11]: import torch

In [12]: import array_api_compat as aac

In [13]: _is_zero_grad = aac.common._helpers._is_jax_zero_gradient_array

In [14]: cfunc = torch.compile(lambda x: _is_zero_grad(x), fullgraph=True)

In [15]: cfunc(torch.ones(3))
Out[15]: False

In [16]: torch.__version__
Out[16]: '2.7.1'

In [17]: aac.__version__
Out[17]: '1.16.0.dev0'

@lucascolley

lucascolley commented Aug 17, 2026

Copy link
Copy Markdown
Member Author

Interesting, I'll have to take a closer look.

@lucascolley
lucascolley marked this pull request as draft August 21, 2026 18:32
@lucascolley

lucascolley commented Aug 26, 2026

Copy link
Copy Markdown
Member Author

weird, I can no longer reproduce the failure... maybe I'm misremembering how I triggered it in the first place. I'll close this for now, can return to it if hit again

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@ev-br ev-br Awaiting requested review from ev-br

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@lucascolley @ev-br